import numpy as np


def get_positional_embeddings(sequence_length, embed_length):
    result = np.zeros((sequence_length, embed_length), dtype=np.float32)
    for pos in range(sequence_length):
        for i in range(embed_length):
            if i % 2 == 0:
                result[pos][i] = np.sin(pos / (10000 ** (i / embed_length)))
            else:
                result[pos][i] = np.cos(pos / (10000 ** ((i - 1) / embed_length)))
    return result


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    plt.imshow(get_positional_embeddings(100, 512), cmap="hot", interpolation="nearest")
    plt.show()
